import json

import requests


def get_token(api_key, secret_key):
    url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=" + api_key + "&client_secret=" + secret_key

    payload = json.dumps("")
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    access_token = response.json().get("access_token")
    model = "ernie-4.0-turbo-8k"
    # model = "ernie-4.0-8k-0329"
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/" + model + "?access_token=" + str(
        access_token)
    return url


def get_answer(url, question):
    data = json.dumps({
        "max_output_tokens": 20,
        "system": "在选项中找出问题的正确答案，请保证信息的准确，并在回答时保持简洁，不需要任何其他分析反馈。",
        "messages": [
            {
                "role": "user",
                "content": question
            }
        ],
        "stream": True
    })
    header = {'Content-Type': 'application/json'}
    resp = requests.request("POST", url, headers=header, data=data, stream=True)
    lines = resp.iter_lines()
    try:
        answer = '│ BD：'
        # print('│ BD：', end='')
        for line in lines:
            if line:
                # print('', end='')
                data_str = line.decode('utf-8')
                if len(data_str) >= 4 and data_str[:4] == 'data':
                    data_str = data_str.split(':', 1)[1].strip()
                    data_dict = json.loads(data_str)
                    # print(data_dict['result'], end='')
                    answer += data_dict['result']
        print(answer)
    except Exception as e:
        print('[报错]' + str(e))


# get_answer(get_token('VD5tmbkMQvCtH9koisgnAiZs', 'Wauw00thAqvLbJ7vBCORAMinXwY8E8KB'),"问题（《本草纲目》的作者是？）选项（周杰伦 李时珍 司马迁 司马光）")
